﻿#include <algorithm>
#include <gtest/gtest.h>
#include <iostream>
#include <vector>

#include <learn/alg/alg_iterator.h>
#include <learn/alg/alg_pia_curve.h>
#include <learn/alg/alg_pia_surface.h>
#include <learn/utils/output.h>

using namespace xi;
using namespace xi::alg;
using namespace xi::geom;

TEST(TestIterator, curvePIA)
{
    {
        Alg_Curve_PIA<3> curvePIA;
        curvePIA.data.push_back(Vec3{0.0, 0.0, 0.0});
        curvePIA.data.push_back(Vec3{0.0, 1.0, 0.0});
        curvePIA.data.push_back(Vec3{1.0, 1.0, 0.0});
        curvePIA.data.push_back(Vec3{1.0, 0.0, 0.0});
        auto func = curvePIA.generate();
        int count = 0;
        int n = 1000;
        for (auto x : func(n))
        {
            count++;
        }
        std::cout << "B-spline PIA count: " << count << std::endl;
    }

    {
        Alg_Curve_PIA<4> curvePIA;
        curvePIA.data.push_back(Vec4{0.0, 0.0, 0.0, 1.0});
        curvePIA.data.push_back(Vec4{0.0, 1.0, 0.0, 0.5});
        curvePIA.data.push_back(Vec4{1.0, 1.0, 0.0, 0.5});
        curvePIA.data.push_back(Vec4{1.0, 0.0, 0.0, 1.0});
        auto func = curvePIA.generate();
        int count = 0;
        int n = 1000;
        for (auto x : func(n))
        {
            count++;
        }
        std::cout << "NURBS PIA count: " << count << std::endl;
    }
}

TEST(TestIterator, curveLSPIA)
{
    {
        Alg_Curve_LSPIA<3> curvePIA;

        auto f = [](double x) -> Vec3 { return Vec3{x, x * x, x * x * x}; };
        for (int i = 0; i < 100; i++)
            curvePIA.data.push_back(f(i / 100.0));
        curvePIA.poleNum = 10;

        auto func = curvePIA.generate();
        int count = 0;
        int n = 1000;
        for (auto x : func(n))
        {
            count++;
        }
        std::cout << "B-spline ExPIA count: " << count << std::endl;
        // std::cout << "B-spline ExPIA error: " << curvePIA.error << std::endl;
    }

    {
        Alg_Curve_FairingPIA<3> curvePIA;

        auto f = [](double x) -> Vec3 { return Vec3{x, x * x, x * x * x}; };
        for (int i = 0; i < 100; i++)
            curvePIA.data.push_back(f(i / 100.0));

        curvePIA.poleNum = 10;
        for (int i = 0; i < curvePIA.poleNum; i++)
        {
            curvePIA.wBending.push_back(1e-8);
            curvePIA.wStretch.push_back(1e-1);
        }

        auto func = curvePIA.generate();
        int count = 0;
        int n = 3000;
        std::shared_ptr<Geom_BSplineCurve<3>> curve;
        for (auto x : func(n))
        {
            curve = x;
            count++;
        }
        std::cout << "B-spline Fairing PIA count: " << count << std::endl;

        auto poles = curve->poles();
        Eigen::MatrixXd A(poles.size(), 3);
        for (int i = 0; i < poles.size(); i++)
        {
            A.row(i) << poles[i].x(), poles[i].y(), poles[i].z();
        }

        auto B0 = Alg_CoeffMatrix::curve_least_square(curvePIA.params, curve->knots(), curve->degree());
        auto [B1, B2] = Alg_CoeffMatrix::curve_stretch_bending(curve->knots(), curve->degree());
        auto Fitting = A.transpose() * B0.transpose() * B0 * A;
        auto Stretching = A.transpose() * B1 * A;
        auto Bending = A.transpose() * B2 * A;
        std::cout << Fitting.trace() << std::endl;
        std::cout << Stretching.trace() << std::endl;
        std::cout << Bending.trace() << std::endl;
    }
}

TEST(TestIterator, surfacePIA)
{
    auto f = [](double u, double v) -> Vec3 { return Vec3{u, v, sin(u + v) * cos(u - v)}; };

    {
        // 20 x 30 个采样点
        int N = 10;
        int M = 12;
        std::vector<std::vector<Vec3>> points(N, std::vector<Vec3>(M));
        for (int i = 0; i < N; i++)
            for (int j = 0; j < M; j++)
                points[i][j] = f(5.0 / (N - 1) * i, 5.0 / (M - 1) * j);

        Alg_Surface_PIA<3> surfacePIA;
        surfacePIA.uDegree = 3;
        surfacePIA.vDegree = 2;
        surfacePIA.data = points;
        auto func = surfacePIA.generate();
        int count = 0;
        int n = 30;
        for (auto x : func(n))
        {
            count++;
        }
        std::cout << "B-spline surface PIA count: " << count << std::endl;
    }

    {
        // 20 x 30 个采样点
        int N = 20;
        int M = 30;

        std::vector<Vec3> points(N * M);
        for (int i = 0; i < N; i++)
            for (int j = 0; j < M; j++)
                points[i * M + j] = f(5.0 / (N - 1) * i, 5.0 / (M - 1) * j);

        Alg_Surface_LSPIA<3> surfacePIA;
        surfacePIA.uDegree = 3;
        surfacePIA.vDegree = 2;
        surfacePIA.uPoleNum = 10;
        surfacePIA.vPoleNum = 8;
        surfacePIA.data = points;
        auto func = surfacePIA.generate();
        int count = 0;
        int n = 30;
        for (auto x : func(n))
        {
            count++;
        }
        std::cout << "B-spline surface LSPIA count: " << count << std::endl;
    }

    {
        // 20 x 30 个采样点
        int N = 20;
        int M = 30;

        std::vector<Vec3> points(N * M);
        for (int i = 0; i < N; i++)
            for (int j = 0; j < M; j++)
                points[i * M + j] = f(5.0 / (N - 1) * i, 5.0 / (M - 1) * j);

        Alg_Surface_FairingPIA<3> surfacePIA;
        surfacePIA.uDegree = 3;
        surfacePIA.vDegree = 2;
        surfacePIA.uPoleNum = 10;
        surfacePIA.vPoleNum = 8;
        for (int i = 0; i < surfacePIA.uPoleNum; i++)
        {
            for (int j = 0; j < surfacePIA.vPoleNum; j++)
            {
                // bending 权设为 1e-1 时，收敛非常慢，迭代 100000 次才能得到平面
                surfacePIA.wBending.push_back(1e-8);
                surfacePIA.wStretch.push_back(1e-8);
            }
        }

        surfacePIA.data = points;
        auto func = surfacePIA.generate();
        int count = 0;
        int n = 30;
        std::shared_ptr<Geom_BSplineSurface<3>> surface;
        for (auto x : func(n))
        {
            surface = x;
        }

        auto poles = surface->poles();
        Eigen::MatrixXd A(poles.size(), 3);
        for (int i = 0; i < poles.size(); i++)
        {
            A.row(i) << poles[i].x(), poles[i].y(), poles[i].z();
        }

        auto B0 = Alg_CoeffMatrix::surface_least_square(surfacePIA.params, surface->u_knots(), surface->v_knots(),
                                                        surfacePIA.uDegree, surfacePIA.vDegree);
        auto [B1, B2] = Alg_CoeffMatrix::surface_stretch_bending(surface->u_knots(), surface->v_knots(),
                                                                 surfacePIA.uDegree, surfacePIA.vDegree);
        auto Fitting = A.transpose() * B0.transpose() * B0 * A;
        auto Stretching = A.transpose() * B1 * A;
        auto Bending = A.transpose() * B2 * A;
        std::cout << Fitting.trace() << std::endl;
        std::cout << Stretching.trace() << std::endl;
        std::cout << Bending.trace() << std::endl;
    }
}